""" 二阶系统镇定问题的常用思路，与技巧
"""

from scipy.integrate import solve_ivp
import numpy as np
import matplotlib.pyplot as plt

class DoubleIntegratorPD:
    ode_method="RK45"
    def __init__(self,kp,kd,delta):
        self.k1=kp
        self.k2=kd
        self.delta=delta
    def control(self,x,dxdt,debug):
        x1,x2=x[0],x[1]
        u=-self.k1*x1-self.k2*x2
        debug['u']=u
        return u

    def solve(self,t_span,initial_state):
        def rhs(t,x,debug={}):
            dxdt=np.zeros_like(x)
            u=self.control(x,dxdt,debug)
            delta=self.delta(t)
            dxdt[0]=x[1]
            dxdt[1]=u+delta
            debug['delta']=delta
            return dxdt

        sol=solve_ivp(rhs,t_span,initial_state,
            dense_output=True,
            t_eval=np.linspace(t_span[0],t_span[-1],1000),
            method=self.ode_method)
        signals=[]
        for it,t in enumerate(sol.t):
            s={}
            rhs(t,sol.y[:,it],debug=s)
            signals.append(s)
        return sol,signals

class DoubleIntegratorPID(DoubleIntegratorPD):
    def __init__(self,kp,ki,kd,delta):
        self.kp=kp
        self.ki=ki
        self.kd=kd
        self.delta=delta
    def control(self,x,dxdt,debug):
        x1,x2,x_int=x[0],x[1],x[2]
        dxdt[2]=x1
        u=-self.kp*x1-self.kd*x2-self.ki*x_int
        debug['u']=u
        return u

class DoubleIntegratorSMC(DoubleIntegratorPD):
    def __init__(self,k,c,delta):
        self.k=k
        self.c=c
        self.delta=delta
    def control(self,x,dxdt,debug):
        x1,x2,x_int=x[0],x[1],x[2]
        dxdt[2]=x1

        s=self.c*x1+x2
        u=-self.k*np.sign(s)-self.c*x2
        # in typical SMC control law, we should appednd -c*v here to get a global convergence,
        # but it is omitted here to make the code more readable.
        debug['u']=u
        debug['s']=s
        return u

def signal_generator(type,period,amplitude):
    def sin(t):
        return amplitude*np.sin(2*np.pi*t/period)
    def trapzoid(t):
        return amplitude*(1-np.mod(t,period)/period)

    if type=="sin":
        return sin
    if type=="trapezoid":
        return trapzoid


if __name__=="__main__":
    plt.figure()
    for i,sig in enumerate(["sin","trapezoid"]):
        plt.subplot(2,1,1+i)
        pd=DoubleIntegratorPD(1,5,signal_generator(sig,5,1))
        sol,signals=pd.solve([0,30],np.array([1,-1,0]))
        plt.plot(sol.t,sol.y[1],label="x2")
        plt.plot(sol.t,sol.y[0],label="x1")
        plt.plot(sol.t,[s["delta"] for s in signals],label="delta")
        plt.grid()
        plt.legend()
    plt.show()

    plt.figure()
    for i,sig in enumerate(["sin","trapezoid"]):
        plt.subplot(2,1,1+i)
        smc=DoubleIntegratorSMC(1.2,5,signal_generator(sig,5,1))
        sol,signals=smc.solve([0,60],np.array([1,-1,0]))
        plt.plot(sol.t,sol.y[1],label="x2")
        plt.plot(sol.t,sol.y[0],label="x1")
        plt.plot(sol.t,[s["delta"] for s in signals],label="delta")
        plt.plot(sol.t,[s["s"] for s in signals],label="s")
        plt.grid()
        plt.legend()
    plt.show()

    # simulation of PID controller is commented out ,
    # since it is not working properly.
    # pid=DoubleIntegratorPID(1,0.8,5,lambda t:np.sin(10*t))
    # sol,signals=pid.solve([0,300],np.array([1,-1,0]))
    # plt.figure()
    # plt.subplot(2,1,1)
    # plt.plot(sol.t,sol.y[1],label="x2")
    # plt.plot(sol.t,sol.y[0],label="x1")
    # plt.grid()
    # plt.legend()
    # plt.subplot(2,1,2)
    # plt.plot(sol.t,sol.y[1],label="x2")
    # plt.plot(sol.t,sol.y[0],label="x1")
    # plt.grid()
    # plt.legend()
    # plt.xlim(280,300)
    # plt.ylim(-0.1,0.1)
    # plt.show()



